package com.ruoyi.kafka;

import lombok.extern.java.Log;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;

import java.time.Duration;
import java.util.*;

/**
 * @program: ruoyi
 * @description: KafkaAdminTest
 * @author: zengchen
 * @create: 2020-10-15 16:39
 **/
@Log
public class KafkaConsumerTest {

    public static long partition0Offset = 0;
    public static long partition1Offset = 0;

    public static void main(String[] args) throws Exception {
        // 官网标准示例  自动提交offset
//        autoOffset();

        // 手动提交offset,  enable.auto.commit=false
        // 关键是  auto.offset.reset 参数，
        // 设置 earliest，每次把offset设置到上一次提交之后，可以消费到从上一次提交offset之后至启动consumer之间的消息
        // 启动期间不会重复消费未提交的消息，重新启动 consumer 会重复消费未提交的消息
        // 设置 latest，每次把offset设置到最后，以前错过的消息就无法重新消费到 （测试结果这个是默认值）
        // 设置 none，如果消息组从来没有提交过，报异常，以前有提交，不报异常，作用同 earliest
        codeOffset();

        // 消费指定的topic分区，消息也可以根据分区单独处理
//        consumePartition();

        // 指定 offset 位置消费和精确提交 offset
//        consumeOffset();

        // 消费暂停和启动
        //  1、接收到record信息以后，去令牌桶中拿取令牌
        //  2、如果获取到令牌，则继续业务处理
        //  3、如果获取不到令牌， 则pause等待令牌
        //  4、当令牌桶中的令牌足够， 则将consumer置为resume状态
        //  consumer.pause(Arrays.asList(topicPartition0, topicPartition1));
        //  consumer.resume(Arrays.asList(topicPartition0, topicPartition1));
//        consumeControl();

    }

    // 指定 offset 位置消费和精确提交 offset
    private static void consumeOffset() {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KafkaAdminTest.BOOTSTRAP_SERVERS_CONFIG);
        props.setProperty("group.id", "test01");
//        props.setProperty("enable.auto.commit", "false");
//        props.setProperty("auto.commit.interval.ms", "1000");
        props.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        props.setProperty("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.setProperty("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);

        TopicPartition topicPartition0 = new TopicPartition(KafkaAdminTest.TOPIC_NAME, 0);
        TopicPartition topicPartition1 = new TopicPartition(KafkaAdminTest.TOPIC_NAME, 1);
        consumer.assign(Arrays.asList(topicPartition0, topicPartition1));
        // 消费订阅哪一个Topic或者几个Topic
//        consumer.subscribe(Arrays.asList(KafkaAdminTest.TOPIC_NAME));
        while (true) {
            System.out.println("partition0Offset=" + partition0Offset + ",partition1Offset=" + partition1Offset);
            consumer.seek(topicPartition0, partition0Offset);
            consumer.seek(topicPartition1, partition1Offset);
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(10000));
            for (TopicPartition partition : records.partitions()) {
                System.out.println("partition=" + partition.partition());
                List<ConsumerRecord<String, String>> recordsPartition = records.records(partition);
                for (ConsumerRecord<String, String> record : recordsPartition) {
                    // 想把数据保存到数据库，成功就成功，不成功...
                    // TODO record 2 db
                    System.out.printf("patition = %d , offset = %d, key = %s, value = %s%n",
                            record.partition(), record.offset(), record.key(), record.value());
                    // 如果失败，则回滚， 不要提交offset
                }

                long lastOffset = recordsPartition.get(recordsPartition.size() - 1).offset();
                // 单个partition中的offset，并且进行提交
                Map<TopicPartition, OffsetAndMetadata> offset = new HashMap<>();
                if (0 == partition.partition()) {
                    partition0Offset = lastOffset + 1; // 下一次消费开始位置
                } else if (1 == partition.partition()) {
                    partition1Offset = lastOffset + 1; // 下一次消费开始位置
                }
                System.out.println("submit lastOffset  = " + lastOffset);
                offset.put(partition, new OffsetAndMetadata(lastOffset));
                // 按专区提交offset
                consumer.commitSync(offset);
            }
        }
    }

    // 消费指定的topic分区，消息也可以根据分区单独处理
    private static void consumePartition() {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KafkaAdminTest.BOOTSTRAP_SERVERS_CONFIG);
        props.setProperty("group.id", "test01");
//        props.setProperty("enable.auto.commit", "false");
//        props.setProperty("auto.commit.interval.ms", "1000");
        props.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        props.setProperty("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.setProperty("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);

        TopicPartition topicPartition0 = new TopicPartition(KafkaAdminTest.TOPIC_NAME, 0);
        TopicPartition topicPartition1 = new TopicPartition(KafkaAdminTest.TOPIC_NAME, 1);
        consumer.assign(Arrays.asList(topicPartition0, topicPartition1));
        // 消费订阅哪一个Topic或者几个Topic
//        consumer.subscribe(Arrays.asList(KafkaAdminTest.TOPIC_NAME));
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
            for (TopicPartition partition : records.partitions()) {
                System.out.println("partition=" + partition.partition());
                List<ConsumerRecord<String, String>> recordsPartition = records.records(partition);
                for (ConsumerRecord<String, String> record : recordsPartition) {
                    // 想把数据保存到数据库，成功就成功，不成功...
                    // TODO record 2 db
                    System.out.printf("patition = %d , offset = %d, key = %s, value = %s%n",
                            record.partition(), record.offset(), record.key(), record.value());
                    // 如果失败，则回滚， 不要提交offset
                }

                // 如果成功，手动通知offset提交
                consumer.commitAsync();
            }
        }
    }

    // 手动提交offset
    private static void codeOffset() {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KafkaAdminTest.BOOTSTRAP_SERVERS_CONFIG);
        props.setProperty("group.id", "test01");
//        props.setProperty("enable.auto.commit", "false");
//        props.setProperty("auto.commit.interval.ms", "1000");
//        props.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        props.setProperty("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.setProperty("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);
        // 消费订阅哪一个Topic或者几个Topic
        consumer.subscribe(Arrays.asList(KafkaAdminTest.TOPIC_NAME));
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
            for (ConsumerRecord<String, String> record : records) {
                // 想把数据保存到数据库，成功就成功，不成功...
                // TODO record 2 db
                System.out.printf("patition = %d , offset = %d, key = %s, value = %s%n",
                        record.partition(), record.offset(), record.key(), record.value());
                // 如果失败，则回滚， 不要提交offset
            }

            // 如果成功，手动通知offset提交
            consumer.commitAsync();
        }
    }

    // 官网标准示例  自动提交offset
    private static void autoOffset() {
        Properties props = new Properties();
        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, KafkaAdminTest.BOOTSTRAP_SERVERS_CONFIG);
        props.setProperty("group.id", "test01");
        props.setProperty("enable.auto.commit", "true");
        props.setProperty("auto.commit.interval.ms", "1000");
        props.setProperty("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.setProperty("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);
        // 消费订阅哪一个Topic或者几个Topic
        consumer.subscribe(Arrays.asList(KafkaAdminTest.TOPIC_NAME));
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(10000));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("patition = %d , offset = %d, key = %s, value = %s%n",
                        record.partition(), record.offset(), record.key(), record.value());
            }
        }
    }


}
